import os
import subprocess
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

class GraphRAGRetriever:
    def __init__(self, graphrag_root_path="/root/graphrag-basic/ragtest"):
        """
        初始化GraphRAG检索器
        :param graphrag_root_path: GraphRAG工作目录（需与部署时的ragtest目录一致）
        """
        self.graphrag_root = graphrag_root_path
        # 验证GraphRAG工作目录是否存在
        if not os.path.exists(self.graphrag_root):
            raise FileNotFoundError(f"GraphRAG工作目录不存在：{self.graphrag_root}，请确认部署路径")
        # 验证GraphRAG配置文件是否存在
        self.settings_path = os.path.join(self.graphrag_root, "settings.yaml")
        if not os.path.exists(self.settings_path):
            raise FileNotFoundError(f"GraphRAG配置文件不存在：{self.settings_path}，请确认部署完整性")

    def retrieve_enhanced_knowledge(self, fault_name, query_template=None, timeout=300):
        """
        调用GraphRAG本地查询获取故障增强知识
        :param fault_name: 标准故障名称（用于构建查询语句）
        :param query_template: 查询模板（默认模板已覆盖故障原因、解决方案、关联组件）
        :param timeout: 查询超时时间（秒）
        :return: GraphRAG返回的增强知识文本（空字符串表示无结果）
        """
        # 默认查询模板：覆盖故障核心维度，可根据需求调整
        if not query_template:
            query_template = f"""
            请提供与故障「{fault_name}」相关的以下信息：
            1. 故障的深层技术原因（补充知识图谱未覆盖的细节）
            2. 行业通用的解决方案（区别于现有预定义方案）
            3. 关联的系统组件或服务（可能影响的上游/下游依赖）
            4. 故障排查的关键命令或工具（辅助验证的实操步骤）
            要求：信息准确、简洁，基于知识库内容，不编造内容。
            """
        
        try:
            # 调用GraphRAG本地查询命令
            result = subprocess.run(
                [
                    "graphrag", 
                    "query", 
                    "--root", self.graphrag_root,
                    "--method", "local",
                    "--query", query_template.strip()  # 去除多余空格，避免命令参数错误
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,  # 合并错误输出，便于排查问题
                text=True,
                timeout=timeout,
                cwd=self.graphrag_root  # 切换到GraphRAG工作目录执行命令
            )
            
            # 处理查询结果：过滤无效信息，提取核心内容
            if result.returncode == 0:
                # 清理输出（去除命令行日志，保留查询结果）
                raw_output = result.stdout.strip()
                # 提取GraphRAG返回的核心结果（根据实际输出格式调整，此处适配默认输出）
                if "Query result:" in raw_output:
                    enhanced_knowledge = raw_output.split("Query result:")[-1].strip()
                elif "Result:" in raw_output:
                    enhanced_knowledge = raw_output.split("Result:")[-1].strip()
                else:
                    enhanced_knowledge = raw_output  # 兼容不同版本的输出格式
                
                # 过滤空结果或无效信息
                if enhanced_knowledge and not any(
                    keyword in enhanced_knowledge.lower() 
                    for keyword in ["no result", "not found", "empty"]
                ):
                    print(f"✅ GraphRAG检索到「{fault_name}」的增强知识")
                    return f"【GraphRAG增强知识】\n{enhanced_knowledge}"
                else:
                    print(f"ℹ️ GraphRAG未检索到「{fault_name}」的相关知识")
                    return ""
            else:
                print(f"❌ GraphRAG查询失败（退出码：{result.returncode}）：{result.stdout[:500]}")
                return ""
        
        except subprocess.TimeoutExpired:
            print(f"❌ GraphRAG查询超时（超过{timeout}秒）")
            return ""
        except Exception as e:
            print(f"❌ GraphRAG检索过程出错：{str(e)}")
            return ""